import numpy as np
import math
import pygad

# -------------------------- 1. 基础参数与复用函数 --------------------------
# 定义场景中的关键坐标和参数
target_true = np.array([0, 200, 0])  # 真目标的位置坐标 (x, y, z)
r, h = 7, 10  # 真目标的半径7m、高10m（圆柱形目标）
target_fake = np.array([0, 0, 0])  # 假目标的位置坐标
missile_M1 = np.array([20000, 0, 2000])  # M1导弹的初始位置
missile_speed = 300  # 导弹速度 (m/s)
drone_FY1 = np.array([17800, 0, 1800])  # FY1无人机的初始位置
g = 9.8  # 重力加速度 (m/s²)
smoke_radius = 10  # 烟幕半径 (m)
smoke_sink_speed = 3  # 烟幕下沉速度 (m/s)
smoke_valid_time = 20  # 烟幕有效时间 (s)
drone_speed_range = [70, 140]  # 无人机速度约束 (m/s)
NUM_SMOKE = 3  # 烟幕弹数量

# 计算导弹飞到假目标的时间（关键约束）
missile_distance_to_fake = np.linalg.norm(target_fake - missile_M1)  # 导弹到假目标的欧氏距离
missile_time_to_fake = missile_distance_to_fake / missile_speed  # 导弹到达假目标的时间 (s)


# 定义获取长方体顶点的函数（用于近似圆柱形目标的边界）
def get_cylinder_bounding_box_vertices(target_center, radius, height):
    """
    计算圆柱形目标的边界长方体顶点
    参数:
        target_center: 目标中心点坐标
        radius: 圆柱半径
        height: 圆柱高度
    返回:
        包含8个顶点坐标的列表
    """
    # 计算长方体的边界
    x_min = target_center[0] - radius
    x_max = target_center[0] + radius
    y_min = target_center[1] - radius
    y_max = target_center[1] + radius
    z_min = target_center[2]  # 圆柱底部高度
    z_max = target_center[2] + height  # 圆柱顶部高度

    # 定义长方体的8个顶点
    vertices = [
        np.array([x_min, y_min, z_min]),  # 底面左下
        np.array([x_max, y_min, z_min]),  # 底面右下
        np.array([x_max, y_max, z_min]),  # 底面右上
        np.array([x_min, y_max, z_min]),  # 底面左上
        np.array([x_min, y_min, z_max]),  # 顶面左下
        np.array([x_max, y_min, z_max]),  # 顶面右下
        np.array([x_max, y_max, z_max]),  # 顶面右上
        np.array([x_min, y_max, z_max])  # 顶面左上
    ]
    return vertices


# 预计算目标顶点（用于后续的遮蔽判断）
target_vertices = get_cylinder_bounding_box_vertices(target_true, r, h)


def missile_real_time_position(missile_init_pos, t):
    """
    计算导弹在时间t的位置
    参数:
        missile_init_pos: 导弹初始位置
        t: 时间 (s)
    返回:
        导弹在时间t的位置坐标
    """
    # 导弹到达假目标后不再移动
    if t >= missile_time_to_fake:
        return target_fake
    if t < 0:
        return missile_init_pos

    # 计算导弹飞行方向向量
    dir_vec = target_fake - missile_init_pos
    dir_norm = np.linalg.norm(dir_vec)
    if dir_norm <= 1e-6:  # 避免除以零
        return missile_init_pos

    # 单位方向向量
    dir_unit_vec = dir_vec / dir_norm

    # 计算导弹位置：初始位置 + 速度 × 时间 × 方向
    return missile_init_pos + missile_speed * t * dir_unit_vec


# -------------------------- 2. 烟幕实时位置计算 --------------------------
def get_smoke_explosion_center(drone_init_pos, drone_speed, theta, t_drop, t_detonate):
    """
    计算烟幕弹爆炸中心点坐标
    参数:
        drone_init_pos: 无人机初始位置
        drone_speed: 无人机速度 (m/s)
        theta: 无人机飞行方向角 (弧度)
        t_drop: 烟幕弹投放延迟时间 (s)
        t_detonate: 烟幕弹起爆延迟时间 (s)
    返回:
        烟幕弹爆炸中心点坐标
    """
    # 计算投放点坐标（无人机飞行t_drop时间后的位置）
    drop_x = drone_init_pos[0] + drone_speed * math.cos(theta) * t_drop
    drop_y = drone_init_pos[1] + drone_speed * math.sin(theta) * t_drop
    drop_z = drone_init_pos[2]  # 投放高度

    # 烟幕弹平抛运动：水平方向继续以无人机速度移动
    explosion_x = drop_x + drone_speed * math.cos(theta) * t_detonate
    explosion_y = drop_y + drone_speed * math.sin(theta) * t_detonate

    # 竖直自由下落：计算下落距离
    fall_distance = 0.5 * g * (t_detonate ** 2)
    explosion_z = max(drop_z - fall_distance, 0.0)  # 确保高度不为负

    return np.array([explosion_x, explosion_y, explosion_z])


def smoke_real_time_center(drone_init_pos, drone_speed, theta, t_drop, t_detonate, t):
    """
    计算烟幕在时间t的中心位置
    参数:
        drone_init_pos: 无人机初始位置
        drone_speed: 无人机速度 (m/s)
        theta: 无人机飞行方向角 (弧度)
        t_drop: 烟幕弹投放延迟时间 (s)
        t_detonate: 烟幕弹起爆延迟时间 (s)
        t: 当前时间 (s)
    返回:
        烟幕在时间t的中心位置坐标
    """
    detonate_t = t_drop + t_detonate  # 烟幕弹起爆时间

    # 如果当前时间小于起爆时间，烟幕尚未形成
    if t < detonate_t - 1e-6:
        return np.array([0, 0, 0])

    # 获取爆炸中心点坐标
    explosion_center = get_smoke_explosion_center(drone_init_pos, drone_speed, theta, t_drop, t_detonate)

    # 计算烟幕下沉距离
    sink_distance = smoke_sink_speed * (t - detonate_t)
    real_z = max(explosion_center[2] - sink_distance, 0.0)  # 确保高度不为负

    return np.array([explosion_center[0], explosion_center[1], real_z])


# -------------------------- 3. 遮蔽判断 --------------------------
def is_segment_sphere_intersect(missile_pos, vertex_pos, smoke_center, smoke_radius):
    """
    判断导弹视线（导弹位置到目标顶点的线段）是否与烟幕球体相交
    参数:
        missile_pos: 导弹当前位置
        vertex_pos: 目标顶点位置
        smoke_center: 烟幕中心位置
        smoke_radius: 烟幕半径
    返回:
       布尔值，表示是否相交
    """
    # 计算向量AC（烟幕中心到导弹位置）和AB（目标顶点到导弹位置）
    vec_AC = smoke_center - missile_pos
    vec_AB = vertex_pos - missile_pos

    # 计算AB向量的平方模长
    ab_sq_mag = np.dot(vec_AB, vec_AB)
    if ab_sq_mag < 1e-12:  # 避免数值误差
        return False

    # 计算投影系数k（AC在AB上的投影长度与AB长度的比值）
    k = np.dot(vec_AC, vec_AB) / ab_sq_mag

    if k <= 0.0:
        # 最近点是导弹位置点A
        min_dist = np.linalg.norm(vec_AC)
    elif 0.0 < k < 1.0:
        # 最近点在线段AB上，使用叉积计算最短距离
        cross_prod = np.cross(vec_AC, vec_AB)
        min_dist = np.linalg.norm(cross_prod) / np.sqrt(ab_sq_mag)
    else:
        # 最近点是目标顶点B
        vec_BC = smoke_center - vertex_pos
        min_dist = np.linalg.norm(vec_BC)

    # 判断最短距离是否小于等于烟幕半径
    return min_dist <= smoke_radius + 1e-6  # 添加小容差避免浮点误差


def is_effective_occlusion(t, drone_speed, theta, t_drop, t_detonate):
    """
    判断在时间t烟幕是否有效遮蔽目标
    参数:
        t: 当前时间 (s)
        drone_speed: 无人机速度 (m/s)
        theta: 无人机飞行方向角 (弧度)
        t_drop: 烟幕弹投放延迟时间 (s)
        t_detonate: 烟幕弹起爆延迟时间 (s)
    返回:
        布尔值，表示是否有效遮蔽
    """
    detonate_t = t_drop + t_detonate  # 烟幕弹起爆时间
    smoke_start = detonate_t  # 烟幕开始时间
    smoke_end = detonate_t + smoke_valid_time  # 烟幕结束时间

    # 检查时间是否在烟幕有效期内
    if t < smoke_start - 1e-6 or t > smoke_end + 1e-6:
        return False

    # 获取导弹当前位置
    missile_pos = missile_real_time_position(missile_M1, t)
    # 获取烟幕当前中心位置
    smoke_center = smoke_real_time_center(drone_FY1, drone_speed, theta, t_drop, t_detonate, t)

    # 检查导弹到每个目标顶点的视线是否被烟幕遮蔽
    for vertex in target_vertices:
        if is_segment_sphere_intersect(missile_pos, vertex, smoke_center, smoke_radius):
            return True

    return False


# 适应度函数：联合遮蔽时长
single_durations = [0.0 for _ in range(NUM_SMOKE)]  # 存储每枚烟幕弹的单独遮蔽时长


def fitness_func(ga_instance, solution, solution_idx):
    """
    遗传算法适应度函数：计算多枚烟幕弹的联合遮蔽时长
    参数:
        ga_instance: 遗传算法实例
        solution: 当前个体解
        solution_idx: 当前个体索引
    返回:
        适应度值（联合遮蔽时长）
    """
    global single_durations

    # 解析解向量：将12个参数分成3组，每组4个参数
    params = []
    for i in range(NUM_SMOKE):
        base = i * 4
        drone_speed = solution[base]  # 无人机速度
        theta = solution[base + 1]  # 飞行方向角
        t_drop = solution[base + 2]  # 投放延迟时间
        t_detonate = solution[base + 3]  # 起爆延迟时间
        params.append((drone_speed, theta, t_drop, t_detonate))

    # 检查投放间隔约束：确保烟幕弹投放时间间隔至少1秒
    t_drop_list = [p[2] for p in params]  # 提取所有投放时间
    t_drop_list.sort()  # 排序投放时间

    # 检查相邻投放时间间隔
    for i in range(len(t_drop_list) - 1):
        if t_drop_list[i + 1] - t_drop_list[i] < 1.0:
            single_durations = [0.0 for _ in range(NUM_SMOKE)]
            return 0.0  # 违反约束，适应度为0

    # 统计每枚烟幕弹的单独干扰时长
    single_durations = []
    time_step = 0.01  # 时间步长 (s)，用于数值积分

    for i in range(NUM_SMOKE):
        drone_speed, theta, t_drop, t_detonate = params[i]
        detonate_t = t_drop + t_detonate  # 起爆时间
        smoke_start = detonate_t  # 烟幕开始时间
        smoke_end = detonate_t + smoke_valid_time  # 烟幕结束时间

        total_duration = 0.0  # 总遮蔽时长
        prev_effective = False  # 上一时间点是否有效遮蔽
        t = smoke_start  # 从烟幕开始时间开始检查

        # 遍历烟幕有效期内的每个时间点
        while t <= smoke_end + 1e-6:
            current_effective = is_effective_occlusion(t, drone_speed, theta, t_drop, t_detonate)

            # 检测遮蔽状态的开始和结束
            if current_effective and not prev_effective:
                start_t = t  # 记录遮蔽开始时间
            elif not current_effective and prev_effective:
                total_duration += t - start_t  # 累加遮蔽时长

            prev_effective = current_effective
            t += time_step

        # 处理烟幕结束时仍处于遮蔽状态的情况
        if prev_effective:
            total_duration += smoke_end - start_t

        single_durations.append(round(total_duration, 3))  # 保留3位小数

    # 统计联合遮蔽时长（多枚烟幕弹共同作用）
    detonate_times = [p[2] + p[3] for p in params]  # 所有烟幕弹的起爆时间
    smoke_start = min(detonate_times)  # 最早起爆时间
    smoke_end = max([dt + smoke_valid_time for dt in detonate_times])  # 最晚结束时间

    total_duration = 0.0  # 联合总遮蔽时长
    prev_effective = False  # 上一时间点是否有效遮蔽
    t = smoke_start  # 从最早起爆时间开始检查

    # 遍历整个联合烟幕期内的每个时间点
    while t <= smoke_end + 1e-6:
        current_effective = False

        # 检查任意一枚烟幕弹是否在当前时间点有效遮蔽
        for i in range(NUM_SMOKE):
            current_effective |= is_effective_occlusion(t, *params[i])
            if current_effective:  # 如果已确定有效，提前退出循环
                break

        # 检测联合遮蔽状态的开始和结束
        if current_effective and not prev_effective:
            start_t = t  # 记录联合遮蔽开始时间
        elif not current_effective and prev_effective:
            total_duration += t - start_t  # 累加联合遮蔽时长

        prev_effective = current_effective
        t += time_step

    # 处理烟幕结束时仍处于遮蔽状态的情况
    if prev_effective:
        total_duration += smoke_end - start_t

    return round(total_duration, 3)  # 返回适应度值（联合遮蔽时长）


# 定义遗传算法的基因空间（变量范围）
gene_space = []
for i in range(NUM_SMOKE):
    gene_space += [
        {'low': 70, 'high': 140},  # 无人机速度范围 [70, 140] m/s
        {'low': 0, 'high': 2 * math.pi},  # 飞行方向角范围 [0, 2π] 弧度
        {'low': 0, 'high': 30},  # 投放延迟时间范围 [0, 30] s
        {'low': 0, 'high': 10}  # 起爆延迟时间范围 [0, 10] s
    ]


def on_generation(ga_instance):
    """
    每一代结束时的回调函数，输出当前迭代信息
    """
    # 获取当前种群最大适应度
    best_fitness = np.max(ga_instance.last_generation_fitness)
    print(f"当前迭代次数：{ga_instance.generations_completed}, 最大遮蔽时长：{best_fitness:.2f}s")


def on_crossover(ga_instance, offspring_crossover):
    """
    交叉操作完成后的回调函数，确保投放时间间隔约束
    """
    for candidate in offspring_crossover:
        t_drop_list = [candidate[i * 4 + 2] for i in range(NUM_SMOKE)]  # 提取投放时间
        t_drop_list.sort()  # 排序投放时间

        # 检查并修正投放时间间隔
        for i in range(len(t_drop_list) - 1):
            if t_drop_list[i + 1] - t_drop_list[i] < 1.0:
                candidate[i * 4 + 2] = t_drop_list[i] + 1.0  # 增加间隔至1秒

    return offspring_crossover


def on_mutation(ga_instance, offspring_mutation):
    """
    变异操作完成后的回调函数，确保投放时间间隔约束
    """
    for candidate in offspring_mutation:
        t_drop_list = [candidate[i * 4 + 2] for i in range(NUM_SMOKE)]  # 提取投放时间
        t_drop_list.sort()  # 排序投放时间

        # 检查并修正投放时间间隔
        for i in range(len(t_drop_list) - 1):
            if t_drop_list[i + 1] - t_drop_list[i] < 1.0:
                candidate[i * 4 + 2] = t_drop_list[i] + 1.0  # 增加间隔至1秒

    return offspring_mutation


def generate_valid_population(pop_size, num_genes, gene_space, num_smoke):
    """
    生成满足投放间隔约束的初始种群
    参数:
        pop_size: 种群大小
        num_genes: 基因数量
        gene_space: 基因空间定义
        num_smoke: 烟幕弹数量
    返回:
        有效初始种群
    """
    population = []

    while len(population) < pop_size:
        candidate = []
        # 随机生成每个烟幕弹的参数
        for i in range(num_smoke):
            candidate += [
                np.random.uniform(gene_space[i * 4]['low'], gene_space[i * 4]['high']),  # 速度
                np.random.uniform(gene_space[i * 4 + 1]['low'], gene_space[i * 4 + 1]['high']),  # 方向角
                np.random.uniform(gene_space[i * 4 + 2]['low'], gene_space[i * 4 + 2]['high']),  # 投放延迟
                np.random.uniform(gene_space[i * 4 + 3]['low'], gene_space[i * 4 + 3]['high'])  # 起爆延迟
            ]

        # 检查投放间隔约束
        t_drop_list = [candidate[i * 4 + 2] for i in range(num_smoke)]
        t_drop_list.sort()
        valid = True

        for i in range(len(t_drop_list) - 1):
            if t_drop_list[i + 1] - t_drop_list[i] < 1.0:
                valid = False
                break

        # 如果满足约束，加入种群
        if valid:
            population.append(candidate)

    return np.array(population)


# 遗传算法主程序
if __name__ == "__main__":
    # 生成满足约束的初始种群
    initial_population = generate_valid_population(
        pop_size=50, num_genes=4 * NUM_SMOKE, gene_space=gene_space, num_smoke=NUM_SMOKE
    )

    # 配置遗传算法参数
    ga_instance = pygad.GA(
        num_generations=150,  # 迭代次数
        num_parents_mating=10,  # 每代参与交配的父代数量
        fitness_func=fitness_func,  # 适应度函数
        sol_per_pop=50,  # 种群大小
        num_genes=4 * NUM_SMOKE,  # 变量数量（12个优化参数）
        gene_space=gene_space,  # 变量范围
        parent_selection_type="tournament",  # 选择策略：锦标赛选择
        crossover_type="two_points",  # 交叉策略：两点交叉
        mutation_type="random",  # 变异策略：随机变异
        mutation_percent_genes=5,  # 变异基因比例（5%）
        mutation_num_genes=1,  # 直接指定变异1个基因
        parallel_processing=8,  # 启用并行计算（8个进程）
        on_generation=on_generation,  # 每一代的回调函数
        on_crossover=on_crossover,  # 交叉操作完成后的回调函数
        on_mutation=on_mutation,  # 变异操作完成后的回调函数
        initial_population=initial_population,  # 自定义初始种群
    )

    print("开始多枚烟幕弹联合优化...")
    # 运行遗传算法
    ga_instance.run()

    # 获取最优解
    solution, solution_fitness, solution_idx = ga_instance.best_solution()

    print("\n最优解：")
    # 解析并输出每枚烟幕弹的最优参数
    for i in range(NUM_SMOKE):
        base = i * 4
        drone_speed = solution[base]  # 最优速度
        theta = solution[base + 1]  # 最优方向角（弧度）
        t_drop = solution[base + 2]  # 最优投放延迟
        t_detonate = solution[base + 3]  # 最优起爆延迟

        # 将弧度转换为角度
        theta_deg = (theta * 180 / math.pi) % 360

        # 计算投放点坐标
        drop_x = drone_FY1[0] + drone_speed * math.cos(theta) * t_drop
        drop_y = drone_FY1[1] + drone_speed * math.sin(theta) * t_drop
        drop_z = drone_FY1[2]

        # 计算起爆点坐标
        explosion_point = get_smoke_explosion_center(
            drone_FY1, drone_speed, theta, t_drop, t_detonate
        )

        # 输出烟幕弹参数
        print(
            f"烟幕弹{i + 1}: 速度={drone_speed:.2f}m/s, 方向角={theta_deg:.2f}°, 投放延迟={t_drop:.2f}s, 起爆延迟={t_detonate:.2f}s")
        print(f"  投放点=({drop_x:.2f}, {drop_y:.2f}, {drop_z:.2f})")
        print(f"  起爆点=({explosion_point[0]:.2f}, {explosion_point[1]:.2f}, {explosion_point[2]:.2f})")
        print(f"  有效干扰时长：{single_durations[i]:.2f}s")

    # 输出联合遮蔽时长
    print(f"联合最大遮蔽时长：{solution_fitness:.2f}s")